'''
Author: Yicheng Chen (yicheng-chen@outlook.com)
LastEditTime: 2024-02-19 16:28:19
'''
import rospy
from nav_msgs.msg import Odometry
from mavros_msgs.msg import PositionTarget
import numpy as np
import pyquaternion
from geometry_msgs.msg import PoseStamped
from geometry_msgs.msg import Twist
from sensor_msgs.msg import Imu
import math

def quart_to_rpy(x, y, z, w):
    roll = math.atan2(2 * (w * x + y * z), 1 - 2 * (x * x + y * y))
    pitch = math.asin(2 * (w * y - x * z))
    yaw = math.atan2(2 * (w * z + x * y), 1 - 2 * (z * z + y * y))
    return roll, pitch, yaw

class FakeMavros():
    def __init__(self, node_name="mavros"):
        # Node
        rospy.init_node(node_name, anonymous=False)

        # Parameters
        self.odom_hz = 30
        init_pos = rospy.get_param("~init_pos", "0.0, 0.0, 0.0")
        self.init_pos = [float(x) for x in init_pos.split(",")]
        
        # Subscribers
        # 此处接收仿真模块反馈的位置信息
        self.local_pos_cmd_sub = rospy.Subscriber('/pose_cmd', PositionTarget, self.cmd_cb)

        # self.local_odom_sub = rospy.Subscriber('/robot1/Odometry',Odometry,self.odom_cb)

        self.local_pos_cmd_sub = rospy.Subscriber('/ugv1/pose', PoseStamped, self.pose_cb)

        # self.local_imu_sub = rospy.Subscriber('/forsense/imu_raw',Imu,self.imu_cb)

        # Publishers
        #这个地方发出了当前的无人机位置
        self.odom_pub = rospy.Publisher('mavros/local_position/odom', Odometry, queue_size=10)
        # self.minidrone_pub = rospy.Publisher('mavros/setpoint_raw/local', PositionTarget, queue_size=10)
        self.ugv_pub = rospy.Publisher('/ugv1/cmd_vel', Twist, queue_size=10)

        # init odom
        self.odom = Odometry()
        self.odom.header.frame_id = "odom"
        self.odom.child_frame_id = "base_link"
        # 初始化里程计的位置
        self.odom.pose.pose.position.x = self.init_pos[0]
        self.odom.pose.pose.position.y = self.init_pos[1]
        self.odom.pose.pose.position.z = self.init_pos[2]
        self.odom.pose.pose.orientation.w = 1

        self.odom.twist.twist.linear.x = 0
        self.odom.twist.twist.linear.y = 0
        self.odom.twist.twist.linear.z = 0
        self.odom.twist.twist.angular.x = 0
        self.odom.twist.twist.angular.y = 0
        self.odom.twist.twist.angular.z = 0
        # init cmd
        self.cmd = PositionTarget()
        self.cmd.coordinate_frame = 1
        self.cmd.type_mask = 0b110111111000
        self.cmd.position.x = None
        self.cmd.position.y = None
        self.cmd.position.z = 0
        self.cmd.velocity.x = 0
        self.cmd.velocity.y = 0
        self.cmd.velocity.z = 0
        self.cmd.acceleration_or_force.x = 0
        self.cmd.acceleration_or_force.y = 0
        self.cmd.acceleration_or_force.z = 0

        # init ugv_cmd
        self.ugv_cmd = Twist()
        self.ugv_cmd.linear.x = 0
        self.ugv_cmd.linear.y = 0
        self.ugv_cmd.linear.z = 0
        self.ugv_cmd.angular.x = 0
        self.ugv_cmd.angular.y = 0
        self.ugv_cmd.angular.z = 0

        self.publish_odom()
        # self.publish_cmd()
        self.ugv_publish_cmd()

    # def odom_cb(self,data):
    #     self.odom = data

    def cmd_cb(self,data):
        self.cmd.position.x = data.position.x
        self.cmd.position.y = data.position.y
        self.cmd.yaw = data.yaw


    def pose_cb(self, data):
        self.odom.header.stamp = rospy.Time.now()
        self.odom.pose.pose.position.x = data.pose.position.x
        self.odom.pose.pose.position.y = data.pose.position.y
        self.odom.pose.pose.position.z = data.pose.position.z

        # Set the orientation in the odom message
        self.odom.pose.pose.orientation.w = data.pose.orientation.w
        self.odom.pose.pose.orientation.x = data.pose.orientation.x
        self.odom.pose.pose.orientation.y = data.pose.orientation.y
        self.odom.pose.pose.orientation.z = data.pose.orientation.z

    def publish_odom(self):
        self.odom_pub_timer = rospy.Timer(rospy.Duration(1/self.odom_hz), self.odom_pub_timer_cb)

    # def publish_cmd(self):
    #     self.cmd_pub_timer = rospy.Timer(rospy.Duration(1/self.odom_hz), self.cmd_pub_timer_cb)

    def ugv_publish_cmd(self):
        self.ugv_cmd_pub_timer = rospy.Timer(rospy.Duration(1/self.odom_hz), self.cmd_ugv_pub_timer_cb)

    def odom_pub_timer_cb(self, event):
        self.odom_pub.publish(self.odom)

    # def cmd_pub_timer_cb(self, event):
    #     self.minidrone_pub.publish(self.cmd)
    
    def cmd_ugv_pub_timer_cb(self, event):
        if(self.cmd.position.x != None and self.cmd.position.y != None):
            euler = quart_to_rpy(self.odom.pose.pose.orientation.x,
                        self.odom.pose.pose.orientation.y,
                        self.odom.pose.pose.orientation.z,
                        self.odom.pose.pose.orientation.w)
            
            psi = euler[2]
            error_yaw = self.cmd.yaw - psi
            error_yaw = math.atan2(math.sin(error_yaw), math.cos(error_yaw)) # wrap it into [-pi, pi]
            error_pos = np.array([self.cmd.position.x - self.odom.pose.pose.position.x, self.cmd.position.y - self.odom.pose.pose.position.y])
            # rospy.INFO(error_pos)
            error_x = np.dot(error_pos, np.array([math.cos(psi),math.sin(psi)])) # pos x error in ugv's frame
            error_y = np.dot(error_pos, np.array([-math.sin(psi),math.cos(psi)])) # pos y error in ugv's frame
            v0 = 0.3
                
            self.ugv_cmd.angular.z = 1.0 * (1.0*error_yaw + 1.0*error_y / math.sqrt(1+pow(np.linalg.norm(error_pos),2)))
            self.ugv_cmd.linear.x = v0*error_yaw + 10.0 * error_x

            self.ugv_pub.publish(self.ugv_cmd)


if __name__ == "__main__":
    mavros = FakeMavros()

    rospy.spin()
